import tensorflow as tf
import argparse
import os

def build_model(num_classes, mode="single"):
    base = tf.keras.applications.EfficientNetV2S(
        include_top=False,
        weights="imagenet",
        pooling="avg",
        input_shape=(384, 384, 3)
    )

    x = tf.keras.layers.Dense(256, activation="relu")(base.output)
    x = tf.keras.layers.Dropout(0.4)(x)

    if mode == "single":
        out = tf.keras.layers.Dense(num_classes, activation="softmax")(x)
        loss = "categorical_crossentropy"
    else:
        out = tf.keras.layers.Dense(num_classes, activation="sigmoid")(x)
        loss = "binary_crossentropy"

    model = tf.keras.Model(inputs=base.input, outputs=out)
    model.compile(optimizer=tf.keras.optimizers.Adam(1e-4), loss=loss, metrics=["accuracy"])
    return model


def main(args):
    train_ds = tf.keras.utils.image_dataset_from_directory(
        args.train_dir,
        image_size=(384, 384),
        batch_size=args.batch_size,
        label_mode='categorical' if args.mode == "single" else 'binary'
    )

    val_ds = tf.keras.utils.image_dataset_from_directory(
        args.val_dir,
        image_size=(384, 384),
        batch_size=args.batch_size,
        label_mode='categorical' if args.mode == "single" else 'binary'
    )

    train_ds = train_ds.map(lambda x, y: (x/255.0, y))
    val_ds = val_ds.map(lambda x, y: (x/255.0, y))

    model = build_model(args.num_classes, args.mode)

    model.fit(
        train_ds,
        validation_data=val_ds,
        epochs=args.epochs,
        callbacks=[
            tf.keras.callbacks.ModelCheckpoint("ultrasound_effv2.h5", save_best_only=True),
            tf.keras.callbacks.EarlyStopping(patience=5, restore_best_weights=True)
        ]
    )

    model.save("ultrasound_effv2_savedmodel")
    print("✅ 模型训练完成并已保存。")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--train_dir", type=str, required=True)
    parser.add_argument("--val_dir", type=str, required=True)
    parser.add_argument("--num_classes", type=int, required=True)
    parser.add_argument("--mode", choices=["single", "multi"], default="single")
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--batch_size", type=int, default=32)
    args = parser.parse_args()
    main(args)
